package org.stayfool.client;

import ch.qos.logback.classic.Level;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.Channel;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.ChannelPipeline;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.mqtt.*;
import io.netty.handler.ssl.SslHandler;
import io.netty.handler.timeout.IdleStateHandler;
import io.netty.util.CharsetUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.stayfool.client.event.*;
import org.stayfool.client.handler.EventHandler;
import org.stayfool.client.handler.HeartBeatHandler;
import org.stayfool.client.handler.MqttClientHandler;
import org.stayfool.client.session.Session;
import org.stayfool.client.session.SessionManager;
import org.stayfool.client.util.*;
import org.stayfool.client.util.future.ConnectFuture;
import org.stayfool.client.util.future.Future;
import org.stayfool.client.util.future.PublishFuture;
import org.stayfool.client.util.future.SubscribeFuture;
import org.stayfool.client.util.message.FixHeaderUtil;
import org.stayfool.client.util.message.IDUtil;

import javax.net.ssl.*;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.KeyStore;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Created by pactera on 2016/11/30.
 *
 * @author stayfool
 */
public class MqttClient {

    private Logger log = LoggerFactory.getLogger(getClass());

    private static final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private static final AtomicInteger clientCount = new AtomicInteger(0);
    private static volatile Bootstrap shareBoot;
    private final AtomicBoolean isConnect = new AtomicBoolean(false);
    private Bootstrap priBoot;
    private Channel channel;
    private final MqttOption option;
    private Session session;

    public MqttClient(MqttOption option) {

        if (option == null || !option.validate())
            throw new IllegalArgumentException();
        this.option = option;

        ch.qos.logback.classic.Logger logger = (ch.qos.logback.classic.Logger) LoggerFactory.getLogger(Logger.ROOT_LOGGER_NAME);
        logger.setLevel(Level.toLevel(option.loglevel()));
    }

    /**
     * connect to server
     */
    public Future<MqttConnAckMessage> connect() {

        initBoot();
        initChannel();
        return doConnect();
    }

    /**
     * subscribe a top
     *
     * @param topic 主题
     * @param qos   qos
     */
    public Future<MqttSubAckMessage> subscribe(String topic, MqttQoS qos) {

        return subscribe(new String[]{topic}, qos);
    }

    /**
     * subscribe topics
     *
     * @param filters 主题
     * @param qos     qos
     */
    public Future<MqttSubAckMessage> subscribe(String[] filters, MqttQoS qos) {

        checkConnect();

        List<MqttTopicSubscription> topicList = new ArrayList<>();
        for (String topic : filters)
            topicList.add(new MqttTopicSubscription(topic, qos));

        MqttFixedHeader fixedHeader = FixHeaderUtil.from(MqttMessageType.SUBSCRIBE);
        MqttMessageIdVariableHeader variableHeader = MqttMessageIdVariableHeader.from(IDUtil.nextAvailableId(option.clientId()));
        MqttSubscribePayload payload = new MqttSubscribePayload(topicList);
        MqttSubscribeMessage msg = new MqttSubscribeMessage(fixedHeader, variableHeader, payload);

        SessionManager.getSession(option.clientId()).waitingConfirm(msg);
        addCallback(EventType.SUBSCRIBE_COMPLETE, (OneTimeEventCallBack<MqttSubAckMessage>) message -> {
            session.confirmMessage(message.variableHeader().messageId());
        });

        SubscribeFuture future = new SubscribeFuture();
        addCallback(EventType.SUBSCRIBE_COMPLETE, future);
        sendMessage(msg);
        return future;
    }

    /**
     * publish a message to a topic
     *
     * @param topic   主题
     * @param qos     qos
     * @param retain  是否保存
     * @param contect 消息内容
     */
    public Future<MqttPubAckMessage> publish(String topic, MqttQoS qos, boolean retain, String contect) {

        checkConnect();

        MqttFixedHeader fixedHeader = new MqttFixedHeader(
                MqttMessageType.PUBLISH,
                false,
                qos,
                retain,
                0
        );
        MqttPublishVariableHeader variableHeader = new MqttPublishVariableHeader(
                topic,
                IDUtil.nextAvailableId(option.clientId())
        );

        ByteBuf payload = Unpooled.wrappedBuffer(contect.getBytes(CharsetUtil.UTF_8));

        MqttPublishMessage msg = new MqttPublishMessage(fixedHeader, variableHeader, payload);
        msg.retain();

        if (qos.value() > MqttQoS.AT_MOST_ONCE.value()) {
            SessionManager.getSession(option.clientId()).waitingConfirm(msg);
            addCallback(EventType.PUBLIST_COMPLETE, (OneTimeEventCallBack<MqttPubAckMessage>) message -> {
                session.confirmMessage(message.variableHeader().messageId());
            });
        }

        PublishFuture future = new PublishFuture();
        addCallback(EventType.PUBLIST_COMPLETE, future);

        sendMessage(msg);

        return future;
    }

    /**
     * cancel a subscription
     *
     * @param topic 主题
     */
    public void unsubscribe(String topic) {

        unsubscribe(new String[]{topic});
    }

    /**
     * cancel some subscription
     *
     * @param filters 主题
     */
    public void unsubscribe(String[] filters) {

        checkConnect();
        MqttFixedHeader fixedHeader = FixHeaderUtil.from(MqttMessageType.UNSUBSCRIBE);
        MqttMessageIdVariableHeader variableHeader = MqttMessageIdVariableHeader.from(IDUtil.nextAvailableId(option.clientId()));
        MqttUnsubscribePayload payload = new MqttUnsubscribePayload(Arrays.asList(filters));

        MqttUnsubscribeMessage msg = new MqttUnsubscribeMessage(fixedHeader, variableHeader, payload);

        SessionManager.getSession(option.clientId()).waitingConfirm(msg);
        addCallback(EventType.UNSUBSCRIBE_SUCCESS, (OneTimeEventCallBack<MqttUnsubAckMessage>) message -> {
            session.confirmMessage(message.variableHeader().messageId());
        });
        sendMessage(msg);
    }

    /**
     * disconnect with server
     */
    public void disconnect() {

        if (isConnect.get()) {
            sendMessage(new MqttMessage(FixHeaderUtil.from(MqttMessageType.DISCONNECT)));
            SessionManager.removeSession(option.clientId());
            isConnect.set(false);
        }
    }

    /**
     * return the connection status
     *
     * @return isConnect
     */
    public boolean isConnect() {

        return isConnect.get();
    }

    /**
     * get clientId
     *
     * @return clientId
     */
    public String getClientId() {

        return option.clientId();
    }

    /**
     * when something than the listener intrest happens , listener will be call
     *
     * @param callback callback
     * @param type     type {@code EventType}
     */
    public void addCallback(EventType type, EventCallback callback) {

        EventManager.register(new EventKey(type, option.clientId()), callback);
    }

    /**
     * remove callback
     *
     * @param type {@code EventType}
     */
    public void removeCallback(EventType type) {

        EventManager.unregister(new EventKey(type, option.clientId()), null);
    }

    /**
     * 返回client的配置
     *
     * @return {@link MqttOption}
     */
    public MqttOption option() {

        return option;
    }

    public <K, V> Session<K, V> getSession() {
        return session;
    }

    public void close() {

        if (isConnect.get())
            disconnect();
        if (!option.shareBoot() || clientCount.get() <= 1) {
            try {
                priBoot.config().group().shutdownGracefully().sync();
            } catch (InterruptedException e) {
            }
        } else {
            clientCount.decrementAndGet();
        }
    }

    private Future<MqttConnAckMessage> doConnect() {

        MqttFixedHeader fixedHeader = FixHeaderUtil.from(MqttMessageType.CONNECT);
        MqttConnectVariableHeader variableHeader = new MqttConnectVariableHeader(
                MqttVersion.MQTT_3_1_1.protocolName(),
                MqttVersion.MQTT_3_1_1.protocolLevel(),
                option.hasUserInfo(),
                option.hasUserInfo(),
                option.willRetain(),
                option.willQos().value(),
                option.willFlag(),
                option.cleanSession(),
                option.keepAlive()
        );
        MqttConnectPayload payload = new MqttConnectPayload(
                option.clientId(),
                option.willTopic(),
                option.willMessage(),
                option.username(),
                option.password()
        );

        ConnectFuture future = new ConnectFuture(isConnect);
        addCallback(EventType.CONNECT_COMPLETE, future);

        sendMessage(new MqttConnectMessage(fixedHeader, variableHeader, payload));

        addCallback(EventType.DIS_CONNECT, message -> this.isConnect.set(false));
        session = SessionManager.init(option);

        return future;
    }

    private Bootstrap createBoot() {

        NioEventLoopGroup workerGroup = new NioEventLoopGroup();
        Bootstrap boot = new Bootstrap();
        boot.group(workerGroup);
        boot.channel(NioSocketChannel.class);
        boot.option(ChannelOption.SO_KEEPALIVE, true);
        boot.handler(new ChannelInitializer<SocketChannel>() {
            @Override
            public void initChannel(SocketChannel ch) throws Exception {

                ChannelPipeline pipeline = ch.pipeline();
                pipeline.addFirst(new HeartBeatHandler());
                pipeline.addFirst(new IdleStateHandler(0, 0, option.keepAlive()));
                pipeline.addLast(new MqttDecoder());
                pipeline.addLast(MqttEncoder.INSTANCE);
                pipeline.addLast(new MqttClientHandler());
                pipeline.addLast(new EventHandler());

                // 如果配置了SSL相关信息，则加入SslHandler
                initSSL(pipeline);
            }
        });

        return boot;
    }

    private void initBoot() {

        if (option.shareBoot()) {
            lock.readLock().lock();
            if (shareBoot == null) {
                lock.readLock().unlock();
                lock.writeLock().lock();
                if (shareBoot == null) {
                    shareBoot = createBoot();
                }
                priBoot = shareBoot;
                lock.writeLock().unlock();
            } else {
                priBoot = shareBoot;
                lock.readLock().unlock();
            }
            clientCount.incrementAndGet();
        } else {
            priBoot = createBoot();
        }
    }

    private void initChannel() {

        try {
            channel = priBoot.connect(option.host(), option.port()).sync().channel();
            ChannelUtil.clientId(channel, option.clientId());
        } catch (InterruptedException e) {
            log.error("init channel failed", e);
        }
    }

    private void initSSL(ChannelPipeline pipeline) {

        if (!option.hasSslInfo())
            return;

        TrustManagerFactory tmf = null;
        KeyManagerFactory kmf = null;
        InputStream is = null;

        Path keypath = Paths.get(option.keyPath());
        if (keypath.isAbsolute()) {
            try {
                is = Files.newInputStream(Paths.get((option.keyPath())));
            } catch (IOException e) {
                log.error("load keys file failed", e);
            }
        } else {
            String keyPath = option.keyPath().replace("classpath:", "");
            ClassLoader cl = Thread.currentThread().getContextClassLoader();
            is = cl.getResourceAsStream(keyPath);
        }

        if (is == null) {
            log.error("SSL key file not found at : {}", option.keyPath());
            return;
        }

        try {
            KeyStore ks = KeyStore.getInstance(KeyStore.getDefaultType());
            ks.load(is, option.keyPass().toCharArray());
            tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            tmf.init(ks);

            KeyManager[] keyManagers = null;
            if (!option.clientMode()) {
                kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
                kmf.init(ks, option.keyPass().toCharArray());
                keyManagers = kmf.getKeyManagers();
            }

            SSLContext ssl = SSLContext.getInstance(MqttOption.SSL);
            ssl.init(keyManagers, tmf.getTrustManagers(), null);
            SSLEngine engine = ssl.createSSLEngine();
            engine.setUseClientMode(option.clientMode());
            pipeline.addFirst(new SslHandler(engine));
        } catch (Exception e) {
            log.error("init SSL failed, use normal mode", e);
        }
    }

    private void checkConnect() {

        assert isConnect.get();
    }

    private void sendMessage(MqttMessage msg) {

        channel.writeAndFlush(msg);
    }

}
